"""
Copyright (c) 2022 Guangdong University of Technology
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Junjiang Xiang
Created: 2023/8/19
Supported by: National Key Research and Development Program of China

"""
import phot
import matplotlib.pyplot as plt


def POL_Constellation(signal_X):
    Fig = plt.figure(figsize=(15, 6))

    ax1 = Fig.add_subplot(1, 1, 1)

    ax1.grid(True)

    ax1.scatter(np.real(signal_X), np.imag(signal_X), s=1)
    title = "X POL Constellation diagram"
    xlabel = "In-phase [a.u.]"
    ylabel = "Quadrature [a.u.]"

    ax1.set_xlabel(xlabel)
    ax1.set_ylabel(ylabel)
    ax1.set_title(title)

    plt.show()

if __name__ == '__main__':

    """单偏振光收发模块 + 光纤信道"""
    """本代码为程序主函数 本代码主要适用于 QPSK, 8QAM, 16QAM, 32QAM, 64QAM 调制格式的单载波相干背靠背（B2B）信号"""

    phot.config(plot=True, backend="numpy")

    """# 系统参数 #"""
    num_symbols = 2**16                              # 符号数目
    bits_per_symbol = 4                              # 2 for QPSK; 4 for 16QAM; 6 for 64QAM  设置调制格式
    total_baud = 30e9                                # 信号波特率，符号率
    up_sampling_factor = 2                           # 上采样倍数
    sampling_rate = up_sampling_factor * total_baud  # 信号采样率
    Reference_frequency = 193.1e12

    """# 首先产生发射端X/Y双偏振信号"""
    signal_bits = phot.gen_bits_single(num_symbols * bits_per_symbol)

    """# 生成双偏振符号序列 #"""
    signals = phot.Modulation_single(signal_bits, bits_per_symbol)

    prev_symbols = signals

    """# 上采样 #"""
    signals = phot.up_sample_single(signals, up_sampling_factor)

    """# RRC #"""
    RRC_ROLL_OFF = 0.01
    signals = phot.pulse_shaper_single(signals, up_sampling_factor, RRC_ROLL_OFF, total_baud)

    """ DAC #"""
    DAC_Resolution_Bits = 10  # DAC量化bit位数
    DAC_OUTPUT_VPP = 1
    signals = phot.dac_single(signals, DAC_Resolution_Bits, DAC_OUTPUT_VPP)

    """ 光源 """
    Linewidth = 100e3
    output_power_dB = 0
    laser_center_frequency = Reference_frequency
    Optical_tx_laser = phot.laser_cw(num_symbols*up_sampling_factor, laser_center_frequency, Reference_frequency, sampling_rate, output_power_dB, Linewidth)

    """ 调制器 """
    Extinction_ratio_db = 60
    VPI = 5
    VDC = -2.5
    signals = phot.IQModulator_POL_single(signals, VDC, Optical_tx_laser, Extinction_ratio_db, VPI)

    """ 发射端电滤波器 """
    filter_type = 'Gaussian'
    BW_electrical = total_baud * 0.75
    order = 4
    signals = phot.Electrical_filter_POL_single(signals, sampling_rate, filter_type, BW_electrical, order)

    """ OSNR """
    osnr_db = 30  # 设置系统OSNR，也就是光信号功率与噪声功率的比值，此处单位为dB
    signals = phot.gaussian_noise_single(signals, osnr_db, sampling_rate)

    P_dB = 0
    power = pow(10, (P_dB / 10)) * 1e-3
    TX_X = signals
    import numpy as np

    TX_X = np.divide(TX_X, np.sqrt(np.mean(np.square(np.abs(TX_X)))))
    TX_X = TX_X * np.sqrt(power)
    signals = TX_X
    # signals = np.expand_dims(signals, axis=1)

    """ Optical Fiber Channel """
    num_spans = 1 # 多少个 span (每个span经过一次放大器)
    span_length = 80  # 一个 span 的长度 (km)
    delta_z = 1  # 单步步长 (km)
    alpha = 0.2  # dB/km
    beta2 = -21.6676e-24  # s^2/km
    gamma = 1.3  # 1/(W*km)
    npol = 1
    
    Signal_Para, Fiber_para, Amplifer_para, PMD_para = phot.Parameters(num_symbols, bits_per_symbol, total_baud,
                                                                       up_sampling_factor, sampling_rate, npol,
                                                                       num_spans, span_length, delta_z, alpha, beta2,
                                                                       gamma)

    signals, signals_power = phot.Fiber(signals, Signal_Para, Fiber_para, PMD_para, Amplifer_para)
    # '''

    """" 相干接收机本振激光器 """
    LO_Linewidth = 100e3                                                    # 激光器线宽
    frequency_offset = 0.5e9                                                # 频偏
    LO_CENTER_Frequency = Reference_frequency + frequency_offset            # LO中心频率
    LO_power_dB = 20                                                        # power OF LO
    Optical_LO = phot.laser_cw(num_symbols*up_sampling_factor, LO_CENTER_Frequency, Reference_frequency, sampling_rate, LO_power_dB, LO_Linewidth)

    """CD补偿"""
    signals = phot.cd_compensation(input=signals, dispersion=17, Lambda=1550, FiberLength=num_spans * span_length,
                                   samplerate=sampling_rate, s=Signal_Para, f=Fiber_para)
    """" 相干接收机 """
    Responsivity = 1
    signals = phot.Receiver_single(signals, Optical_LO, Responsivity)

    """ 接收端滤波器 """
    filter_type = 'Gaussian'
    BW_electrical = total_baud * 0.75
    order = 4
    signals = phot.Electrical_filter_POL_single(signals, sampling_rate, filter_type, BW_electrical, order)

    """ 模拟接收机造成的I/Q失衡，主要考虑幅度失衡和相位失衡，这里将两者都加在虚部上 """
    signals = phot.add_iq_imbalance_single(signals)

    """ 加入ADC的量化噪声 """
    adc_sample_rate = 120e9  # ADC采样率
    adc_resolution_bits = 8  # ADC的bit位数
    signals = phot.adc_noise_single(signals, sampling_rate, adc_sample_rate, adc_resolution_bits)

    """ GSOP """
    signals = phot.iq_compensation_single(signals)

    """ 粗频偏估计和补偿 """
    signals = phot.freq_offset_compensation_single(signals, sampling_rate)

    signals = phot.pulse_shaper_single(signals, up_sampling_factor, RRC_ROLL_OFF, total_baud)

    # 同步 #
    signals, prev_symbols = phot.synchronization_single(signals, prev_symbols, up_sampling_factor)

    """ 自适应均衡/解偏振：采用CMA-RDE算法 """
    # 注意数据导致的标签正负性
    num_tap = 25  # 均衡器抽头数目，此处均衡器内部是采用FIR滤波器，具体可查阅百度或者论文，
    ref_power_cma = 2  # 设置CMA算法的模
    cma_convergence = 30000  # CMA预均衡收敛的信号长度
    step_size_cma = 1e-9  # CMA的更新步长，梯度下降法的步长
    step_size_rde = 1e-9  # RDE的更新步长，梯度下降法的步长，%% CMA和RDE主要就是损失函数不同
    signals = phot.adaptive_equalize_single(
        signals,
        num_tap,
        cma_convergence,
        ref_power_cma,
        step_size_cma,
        step_size_rde,
        up_sampling_factor,
        bits_per_symbol,
        total_baud,
    )

    # 精确的频偏估计和补偿：采用FFT-FOE算法，
    signals = phot.freq_offset_compensation_single(signals, total_baud)

    """ 相位恢复：采用盲相位搜索算法进行相位估计和补偿 """
    N_Test_Angle = 64       # BPS算法的测试角数目
    Block_Size = 100        # BPS算法的块长设置
    signals = phot.bps_restore_single(signals, N_Test_Angle, Block_Size, bits_per_symbol)

    """ 同步 """
    signals, prev_symbols = phot.synchronization_single(signals, prev_symbols, 1)

    """ 相位校正，可供选择 """
    signals, prev_symbols = phot.Correct_phase_single(prev_symbols, signals)

    # 星座图
    POL_Constellation(signals)
    """ 解码 """
    signals = phot.deModulation_single(signals, bits_per_symbol)
    prev_symbols = phot.deModulation_single(prev_symbols, bits_per_symbol)

    """ 计算BER """
    ber, Q = phot.ber_count(prev_symbols, signals)

    print(ber)
    print(Q)